#include<iostream>
#include<vector>
using namespace std;

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        int m = mat.size();
        int n = mat[0].size();
        int count = m + n - 1;
        int posi = 0;
        int posj = 0;
        int tmpDigit = 0;
        vector<int> ans;
        while (tmpDigit < count) {
            while (tmpDigit < count) {
                if (posi >= 0 && posj < n) {
                    ans.push_back(mat[posi][posj]);
                    posi--;
                    posj++;
                }
                else if (posi < 0 && posj < n) {
                    posi++;
                    break;
                }
                else if (posj >= n) {
                    posi += 2;
                    posj--;
                    break;
                }
            }
            tmpDigit++;
            while (tmpDigit < count) {
                if (posi < m && posj >= 0) {
                    ans.push_back(mat[posi][posj]);
                    posi++;
                    posj--;
                }
                else if (posi < m && posj < 0) {
                    posj++;
                    break;
                }
                else if (posi >= m) {
                    posi--;
                    posj += 2;
                    break;
                }
            }
            tmpDigit++;
        }
        return ans;
    }
};